# coding = utf-8

'''
non_local unet
论文地址:https://arxiv.org/pdf/1812.04103.pdf
参考代码：https://github.com/Whu-wxy/Non-local-U-Nets-2D-block
'''


import torch
import torch.nn as nn
import torch.nn.functional as F

import torch
import torch.nn.functional as F
import torch.nn as nn


class multi_head_attention_2d(torch.nn.Module):
    def __init__(self, in_channel, key_filters, value_filters,
                 output_filters, num_heads, dropout_prob=0.5, layer_type='SAME'):
        super().__init__()
        """Multihead scaled-dot-product attention with input/output transformations.

        Args:
            inputs: a Tensor with shape [batch, h, w, channels]
            key_filters: an integer. Note that queries have the same number 
                of channels as keys
            value_filters: an integer
            output_depth: an integer
            num_heads: an integer dividing key_filters and value_filters
            layer_type: a string, type of this layer -- SAME, DOWN, UP
        Returns:
            A Tensor of shape [batch, _h, _w, output_filters]

        Raises:
            ValueError: if the key_filters or value_filters are not divisible
                by the number of attention heads.
        """

        if key_filters % num_heads != 0:
            raise ValueError("Key depth (%d) must be divisible by the number of "
                             "attention heads (%d)." % (key_filters, num_heads))
        if value_filters % num_heads != 0:
            raise ValueError("Value depth (%d) must be divisible by the number of "
                             "attention heads (%d)." % (value_filters, num_heads))
        if layer_type not in ['SAME', 'DOWN', 'UP']:
            raise ValueError("Layer type (%s) must be one of SAME, "
                             "DOWN, UP." % (layer_type))

        self.num_heads = num_heads
        self.layer_type = layer_type

        self.QueryTransform = None
        if layer_type == 'SAME':
            self.QueryTransform = nn.Conv2d(in_channel, key_filters, kernel_size=1, stride=1,
                                            padding=0, bias=True)
        elif layer_type == 'DOWN':
            self.QueryTransform = nn.Conv2d(in_channel, key_filters, kernel_size=3, stride=2,
                                            padding=1, bias=True)  # author use bias
        elif layer_type == 'UP':
            self.QueryTransform = nn.ConvTranspose2d(in_channel, key_filters, kernel_size=3, stride=2,
                                                     padding=1, bias=True)

        self.KeyTransform = nn.Conv2d(in_channel, key_filters, kernel_size=1, stride=1, padding=0, bias=True)
        self.ValueTransform = nn.Conv2d(in_channel, value_filters, kernel_size=1, stride=1, padding=0, bias=True)
        self.attention_dropout = nn.Dropout(dropout_prob)

        self.outputConv = nn.Conv2d(value_filters, output_filters, kernel_size=1, stride=1, padding=0, bias=True)

        self._scale = (key_filters // num_heads) ** 0.5

    def forward(self, inputs):
        """
        :param inputs: B, C, H, W
        :return: inputs: B, Co, Hq, Wq
        """

        if self.layer_type == 'SAME' or self.layer_type == 'DOWN':
            q = self.QueryTransform(inputs)
        elif self.layer_type == 'UP':
            q = self.QueryTransform(inputs, output_size=(inputs.shape[2] * 2, inputs.shape[3] * 2))

        # [B, Hq, Wq, Ck]
        k = self.KeyTransform(inputs).permute(0, 2, 3, 1)
        v = self.ValueTransform(inputs).permute(0, 2, 3, 1)
        q = q.permute(0, 2, 3, 1)

        Batch, Hq, Wq = q.shape[0], q.shape[1], q.shape[2]

        # [B, H, W, N, Ck]
        k = self.split_heads(k, self.num_heads)
        v = self.split_heads(v, self.num_heads)
        q = self.split_heads(q, self.num_heads)

        # [(B, H, W, N), c]
        k = torch.flatten(k, 0, 3)
        v = torch.flatten(v, 0, 3)
        q = torch.flatten(q, 0, 3)

        # normalize
        q = q / self._scale
        # attention
        # [(B, Hq, Wq, N), (B, H, W, N)]
        A = torch.matmul(q, k.transpose(0, 1))
        A = torch.softmax(A, dim=1)
        A = self.attention_dropout(A)

        # [(B, Hq, Wq, N), C]
        O = torch.matmul(A, v)
        # [B, Hq, Wq, C]
        O = O.view(Batch, Hq, Wq, v.shape[-1] * self.num_heads)
        # [B, C, Hq, Wq]
        O = O.permute(0, 3, 1, 2)
        # [B, Co, Hq, Wq]
        O = self.outputConv(O)

        return O

    def split_heads(self, x, num_heads):
        """Split channels (last dimension) into multiple heads.

        Args:
            x: a Tensor with shape [batch, h, w, channels]
            num_heads: an integer

        Returns:
            a Tensor with shape [batch, h, w, num_heads, channels / num_heads]
        """

        channel_num = x.shape[-1]
        return x.view(x.shape[0], x.shape[1], x.shape[2], num_heads, int(channel_num / num_heads))


class residual_block(nn.Module):
    def __init__(self, input_channel, output_channel, stride, is_short_cut):
        super().__init__()
        self.re_bn1 = nn.Sequential(nn.BatchNorm2d(input_channel), nn.ReLU(inplace=True))
        self.is_shor_cut = is_short_cut
        self.short_cut = nn.Conv2d(input_channel, output_channel, 3, stride=stride, padding=1)
        self.conv1 = nn.Conv2d(input_channel, output_channel, 3, stride=stride, padding=1)
        self.re_bn2 = nn.Sequential(nn.BatchNorm2d(output_channel), nn.ReLU(inplace=True))
        self.conv2 = nn.Conv2d(output_channel, output_channel, 3, padding=1)

    def forward(self, x):
        shortcut = x
        x = self.re_bn1(x)
        if self.is_shor_cut:  ###上下采样阶段，需要变换
            shortcut = self.short_cut(x)
        x = self.conv1(x)
        x = self.re_bn2(x)
        x = self.conv2(x)

        return x + shortcut

class attention_block(nn.Module):
    def __init__(self, input_channel, output_channel, stride, is_short_cut):
        super().__init__()
        self.re_bn1 = nn.Sequential(nn.BatchNorm2d(input_channel), nn.ReLU(inplace=True))
        self.is_shor_cut = is_short_cut
        self.short_cut = nn.ConvTranspose2d(input_channel, output_channel, 2, stride=2)

        if stride != 1:
            layer_type = 'UP'
        else:
            layer_type = 'SAME'

        self.attention = multi_head_attention_2d(in_channel=input_channel, key_filters=output_channel,
                                                  value_filters=output_channel, output_filters=output_channel,
                                                  num_heads=2, layer_type=layer_type)

    def forward(self, input):
        shortcut = input
        x = self.re_bn1(input)
        if self.is_shor_cut:
            shortcut = self.short_cut(x)
        x = self.attention(x)
        return x+shortcut



class encoder_layer(nn.Module):
    def __init__(self, input_channel, output_channel, stride, bolcks):
        super().__init__()
        self.block = residual_block(input_channel, output_channel, stride, is_short_cut=True)
        self.block_list = []
        for i in range(1, bolcks):
            self.block_list.append(residual_block(output_channel, output_channel, stride=1, is_short_cut=False))
        self.block_list = nn.ModuleList(self.block_list)

    def forward(self, input):
        x = self.block(input)
        for item in self.block_list:
            x = item(x)

        return x

class decoder_layer(nn.Module):
    def __init__(self, input_channel, output_channel, stride, bolcks):
        super().__init__()
        self.decov1 = nn.ConvTranspose2d(input_channel, output_channel, stride=stride, kernel_size=2)
        self.block_list = []
        for i in range(1, bolcks):
            self.block_list.append(residual_block(output_channel, output_channel, stride=1, is_short_cut=False))
        self.block_list = nn.ModuleList(self.block_list)

    def forward(self, input, skip):
        x = self.decov1(input)
        x = x + skip
        for item in self.block_list:
            x = item(x)

        return x

class attention_decoder_layer(nn.Module):
    def __init__(self, input_channel, output_channel, stride, bolcks):
        super().__init__()
        self.attention = attention_block(input_channel, output_channel, stride, is_short_cut=True)
        self.block_list = []
        for i in range(1, bolcks):
            self.block_list.append(residual_block(output_channel, output_channel, stride=1, is_short_cut=False))

    def forward(self, input, skip):
        x = self.attention(input)
        x = x + skip
        for item in self.block_list:
            x = item(x)

        return x


class NonLocalUnet(nn.Module):
    def __init__(self, input_channel, output_channel, num_channel=32):
        super().__init__()

        self.conv1 = nn.Conv2d(input_channel, num_channel, 3, padding=1)
        self.encode1 = encoder_layer(num_channel, num_channel, stride=1, bolcks=1)
        self.encode2 = encoder_layer(num_channel, num_channel*2, stride=2, bolcks=2)
        self.encode3 = encoder_layer(num_channel*2, num_channel*4, stride=2, bolcks=2)

        self.re_bn1 = nn.Sequential(nn.BatchNorm2d(num_channel*4), nn.ReLU(inplace=True))

        self.bottom_block = multi_head_attention_2d(in_channel=num_channel*4, key_filters=num_channel*4,
                                                  value_filters=num_channel*4, output_filters=num_channel*4,
                                                  num_heads=2, layer_type="SAME")

        self.decode3 = decoder_layer(num_channel*4, num_channel*2, stride=2, bolcks=2)
        self.decode2 = attention_decoder_layer(num_channel*2, num_channel, stride=2, bolcks=2)
        self.output = nn.Conv2d(num_channel, output_channel, 1)

    def forward(self, input):
        x = self.conv1(input)
        e1 = self.encode1(x)
        e2 = self.encode2(e1)
        x = self.encode3(e2)
        x = self.re_bn1(x)
        x = self.bottom_block(x)
        x = self.decode3(x, e2)
        x = self.decode2(x, e1)
        x = self.output(x)
        return x










if __name__ == '__main__':
    unet = NonLocalUnet(input_channel=1, output_channel=2).cuda()
    from torchsummary import summary

    summary(unet, (1, 512, 512))

